Hmmn - this adds some code complexity that may complicate maintenance for a relatively niche feature. Probably the easiest way to extract these without huge memory is just to make a function that extracts the layers you want and then compile it? Torch should be able to correctly drop the unwanted tensors and avoid OOM in that scenario.

Might be useful for certain use-cases as in #33698 ig. Depends on how much memory is being wasted on VLMs/diffusers, it's a common pattern for them to output hidden states

I guess, but I think we should try to avoid adding features that need updates across all modeling code unless they're really important. Plus I think the compile() solution should work fine!

Yeah, we defi don't want to update all models. Anyway, just wanted to link related issues to have a full picture :)

@Rocketknight1 this is actually a common pattern for this feature.

Unless you want to support all models, you don't need to update all modeling. Only some legacy models are using per modeling pattern

Good point, I forgot we'd factored that out into _can_record_outputs for a lot of models. If you want to attempt a PR to update the generic code there to allow specific layer selection, I guess we could consider it! No guarantees, though, I'd need to talk with the rest of the team about whether the extra complexity is worth it

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=43213&sha=25a8fa

@Rocketknight1 could you take a look when you get a chance? This is waiting on a maintainer: team decision on whether layer-selective recording is worth the complexity (Rocketknight1: 'No guarantees... need to talk with the rest of the team').

Thanks for the PR! The output-recording code this edits was rewritten and moved to src/transformers/utils/output_capturing.py in #43765, so this diff no longer applies (none of the changed lines exist on main). If selecting specific layers via output_hidden_states=[...] is still something you want, could you re-implement it on top of the new output_capturing.py hooks? Otherwise we can close this. Closing this one for now; feel free to open a fresh PR against main if it's still needed. Thanks!